import os
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import pyarrow as pa
from dora import Node, DoraStatus
from mofa.kernel.utils.util import create_agent_output
from mofa.utils.files.dir import get_relative_path


def init_browser():
    """初始化Chrome浏览器"""
    chrome_options = Options()
    chrome_options.add_argument("--headless")  # 无头模式
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--no-sandbox")
    try:
        service = Service('/opt/homebrew/bin/chromedriver')  # 替换为你的 ChromeDriver 路径
        browser = webdriver.Chrome(service=service, options=chrome_options)
    except Exception as e:
        print(f"Error starting ChromeDriver: {e}")
        return None
    return browser


def search_bing(query):
    """在Bing中执行搜索并返回结果"""
    browser = init_browser()
    browser.get('https://cn.bing.com')

    search_box = browser.find_element(By.NAME, 'q')
    search_box.send_keys(query)
    search_box.send_keys(Keys.RETURN)

    # 等待结果加载
    WebDriverWait(browser, 10).until(
        EC.presence_of_element_located((By.ID, 'b_results'))
    )

    results = []
    search_results = browser.find_elements(By.CSS_SELECTOR, 'li.b_algo')
    for result in search_results:
        try:
            title_element = result.find_element(By.TAG_NAME, 'h2')
            link_element = result.find_element(By.TAG_NAME, 'a')
            description_element = result.find_element(By.CLASS_NAME, 'b_caption')

            title = title_element.text
            link = link_element.get_attribute('href')
            description = description_element.text

            results.append({
                'title': title,
                'link': link,
                'description': description
            })
        except Exception as e:
            print(f"Error parsing result: {e}")

    browser.quit()
    print(f"Search results: {results}")
    return results


class Operator:
    def on_event(self, dora_event, send_output) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            task = dora_event["value"][0].as_py()  # 获取搜索关键词
            search_results = search_bing(task)  # 在Bing上搜索

            # 将搜索结果转化为 agent_output 格式
            output_data = [{'title': res['title'], 'link': res['link'], 'description': res['description']} for res in search_results]

            # 发送搜索结果
            send_output("search_results", pa.array([create_agent_output(
                step_name='search_results',
                output_data=output_data,
                dataflow_status=os.getenv('IS_DATAFLOW_END', False)
            )]), dora_event['metadata'])

        return DoraStatus.CONTINUE


if __name__ == "__main__":
    query = "你好"
    search_results = search_bing(query)
    for result in search_results:
        print(result)